#include <Windows.h>
#include <Shlwapi.h.>
#include<TlHelp32.h> 
#include <Psapi.h>
#include <tchar.h>
#include "MyUtil.h"
#pragma comment(lib,"Psapi.lib")
#pragma comment(lib,"Shlwapi.lib")

std::vector<MyUtil::ProcessUtil::ProcessInfo> MyUtil::ProcessUtil::processes;
const std::vector<MyUtil::ProcessUtil::ProcessInfo>* MyUtil::ProcessUtil::GetProcessesInfo()
{
	PROCESSENTRY32 pe32;
	pe32.dwSize = sizeof(PROCESSENTRY32);
	BOOL bRet;
	HANDLE hProcessSnap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, NULL);
	bRet = Process32First(hProcessSnap, &pe32);
	DWORD dwCurrentPId = GetCurrentProcessId();
	while (bRet)
	{
		HMODULE hMods[MyUtil::ProcessUtil::nMaxModules] = {0};
		DWORD cbNeeded = 0;
		TCHAR lpstrModName[MAX_PATH];
		BOOL bWow64Process;
		if(0 ==  pe32.th32ProcessID)
		{
			DWORD dwLastError = GetLastError();
			bRet = Process32Next(hProcessSnap,&pe32);
			continue;
		}
		HANDLE hProcess = ::OpenProcess(
			PROCESS_ALL_ACCESS,
			FALSE, pe32.th32ProcessID);
		if(NULL == hProcess)
		{
			DWORD dwLastError = GetLastError();
			bRet = Process32Next(hProcessSnap,&pe32);
			continue;
		}
		IsWow64Process(hProcess, &bWow64Process);
		EnumProcessModulesEx(hProcess, hMods, sizeof(hMods), &cbNeeded, bWow64Process?LIST_MODULES_64BIT:LIST_MODULES_32BIT);
		MyUtil::ProcessUtil::ProcessInfo processInfo;
		processInfo.dwProcessId  = pe32.th32ProcessID;
		std::vector<MyUtil::ProcessUtil::ModuleInfo> modules;
		for (UINT i = 0; i < (cbNeeded / sizeof(HMODULE)); i++ )
		{
			GetModuleFileNameEx(hProcess, hMods[i], lpstrModName, _countof(lpstrModName));
			MODULEINFO modInfo = {0};
			GetModuleInformation(hProcess, hMods[i], &modInfo, sizeof(modInfo));
			//
			MyUtil::ProcessUtil::ModuleInfo moduleInfo;
			moduleInfo.dwImageBase = (DWORD)hMods[i];
			moduleInfo.dwImageSize = modInfo.SizeOfImage;
			lstrcpy(moduleInfo.lpstrPath, lpstrModName);
			LPSTR lpstrExtName = PathFindExtension(lpstrModName);
			LPSTR lpstrFileName = PathFindFileName(lpstrModName);
			lstrcpy(moduleInfo.lpstrFileName, lpstrFileName);
			!lstrcmp(lpstrExtName, _T(".exe")) ? moduleInfo.eType = E_EXE:moduleInfo.eType = E_DLL;
			modules.push_back(moduleInfo);
		}
		if(0 != modules.size())
		{
			processInfo.modules = modules;
			processes.push_back(processInfo);
		}
		CloseHandle(hProcess);
		bRet = Process32Next(hProcessSnap,&pe32);
	}
	return &processes;
}

bool MyUtil::ProcessUtil::SetPrivilege(LPCSTR  lpstrName, bool bEnablePrivilege)
{
	HANDLE hToken;
	BOOL bOpenFlag = OpenProcessToken( GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES, &hToken );
	if( !bOpenFlag )
	{
		return false;
	}
	else 
	{
		TOKEN_PRIVILEGES tp;
		LUID luid;
		if ( !LookupPrivilegeValue( NULL, SE_DEBUG_NAME, &luid ) ) 
		{
			return false;
		}
		tp.PrivilegeCount = 1;
		tp.Privileges[0].Luid = luid;
		if ( bEnablePrivilege ) 
		{
			tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
		}
		else 
		{
			tp.Privileges[0].Attributes = 0;
		}
		if ( !AdjustTokenPrivileges(
			hToken, 
			FALSE, 
			&tp, 
			sizeof(TOKEN_PRIVILEGES), 
			(PTOKEN_PRIVILEGES) NULL, 
			(PDWORD) NULL) )
		{ 
			return false;
		} 
		if (GetLastError() == ERROR_NOT_ALL_ASSIGNED)
		{
			return false;
		} 
	}
	CloseHandle(hToken);
	return true;
}


#define OUT_PUT_BUFF_SIZE 4097
void  MyUtil::OutputDebugStringF(const char *fmt, ...)
{
	va_list vlArgs;
	char *strBuf = (char*)GlobalAlloc(GPTR, OUT_PUT_BUFF_SIZE);
	ZeroMemory(strBuf, OUT_PUT_BUFF_SIZE);
	va_start(vlArgs, fmt);
	_vsnprintf_s(strBuf, OUT_PUT_BUFF_SIZE - 1, OUT_PUT_BUFF_SIZE -1, fmt, vlArgs);
	va_end(vlArgs);
	strcat_s(strBuf, OUT_PUT_BUFF_SIZE - 1, "\n");
	OutputDebugStringA(strBuf);
	GlobalFree(strBuf);
}

void MyUtil::Nothing(const char *fmt, ...){}
